import numpy as np
# 队列
from collections import deque


# Skip-gram + 负采样的数据处理，将数据整理成神经网络需要的类型
class InputData:
    # 初始化参数
    def __init__(self, input_file_name, min_count):
        # 文件名
        self.input_file_name = input_file_name
        # 初始化索引
        self.index = 0
        # 打开文件夹
        self.input_file = open(self.input_file_name, "r", encoding="utf-8")
        # 控制词表的大小，删除词频低于min_count的词
        self.min_count = min_count
        # 初始化词频率字典，key是id，值是频率，用于负采样
        self.wordid_frequency_dict = dict()
        self.word_count = 0
        # 统计单词一共有多少个
        self.word_count_sum = 0
        # 统计句子一共有多少个
        self.sentence_count = 0
        # 初始化id2word字典
        self.id2word_dict = dict()
        # 初始化word2id字典
        self.word2id_dict = dict()
        # 初始化字典
        self._init_dict()
        self.sample_table = []
        # 初始化负采样映射表
        self._init_sample_table()
        # 得到单词索引列表
        self.get_wordId_list()
        # 初始化队列
        self.word_pairs_queue = deque()
        # 结果展示
        print('Word Count is:', self.word_count)
        print('Word Count Sum is', self.word_count_sum)
        print('Sentence Count is:', self.sentence_count)

    def _init_dict(self):
        """
        初始化字典
        """
        # 初始化词频率字典，key是单词，值是频率
        word_freq = dict()
        # 遍历文件，文件其实只有一行
        for line in self.input_file:
            # 删除头尾空格并按照空格进行分割
            line = line.strip().split()
            # 单词总数加上每一行的长度
            self.word_count_sum += len(line)
            # 句子总数加1
            self.sentence_count += 1
            # 遍历每一个单词
            for i, word in enumerate(line):
                # 每隔1000000个单词打印句长
                if i % 1000000 == 0:
                    print(i, len(line))
                # 构建词频率字典
                if word_freq.get(word) == None:
                    word_freq[word] = 1
                else:
                    word_freq[word] += 1

        # 遍历词频率字典
        for i, word in enumerate(word_freq):
            # 每隔100000个单词打印词频率字典的长度
            if i % 100000 == 0:
                print(i, len(word_freq))
            # 如果单词的词频率小于最小频率，则单词总数中减去该单词的数目
            if word_freq[word] < self.min_count:
                self.word_count_sum -= word_freq[word]
                continue
            # 构建word2id字典
            self.word2id_dict[word] = len(self.word2id_dict)
            # 构建id2word字典
            self.id2word_dict[len(self.id2word_dict)] = word
            # 构建词频率字典，key是id,value是词频
            self.wordid_frequency_dict[len(self.word2id_dict) - 1] = word_freq[word]
        # 记录去重单词总数
        self.word_count = len(self.word2id_dict)

    def _init_sample_table(self):
        """
        初始化负采样映射表，生成负样本采样频率
        """
        # 设定词表大小
        sample_table_size = 1e8
        # 根据负采样的公式，先对频率做0.75次方
        pow_frequency = np.array(list(self.wordid_frequency_dict.values())) ** 0.75
        # 公式中的Z
        word_pow_sum = sum(pow_frequency)
        # 做归一化
        ratio_array = pow_frequency / word_pow_sum
        # 统计在1e8大小词表中出现几次
        word_count_list = np.round(ratio_array * sample_table_size)
        # 构建负采样映射表
        for word_index, word_freq in enumerate(word_count_list):
            self.sample_table += [word_index] * int(word_freq)
        self.sample_table = np.array(self.sample_table)
        # 将负采样映射表打乱
        np.random.shuffle(self.sample_table)

    def get_wordId_list(self):
        """
        得到单词索引列表
        """
        # 读取文件
        self.input_file = open(self.input_file_name, encoding="utf-8")
        # 从文件中读取整行
        sentence = self.input_file.readline()
        # 一句中的所有word 对应的 id
        wordId_list = []
        # 整行句子去除首尾空格并以空格进行分割
        sentence = sentence.strip().split(' ')
        # 遍历每一个单词
        for i, word in enumerate(sentence):
            # 每隔1000000个单词，打印句长
            if i % 1000000 == 0:
                print(i, len(sentence))

            try:
                # 从word2id字典中映射得到每个词的索引
                word_id = self.word2id_dict[word]
                # 将单词索引添加到wordId_list
                wordId_list.append(word_id)
            except:
                continue
        # 得到单词索引列表
        self.wordId_list = wordId_list

    def get_batch_pairs(self, batch_size, window_size):
        """
        正采样
        :param batch_size: batch大小
        :param window_size: 窗口大小
        :return:
        """
        # 当单词队列长度小于batch size时，则构建单词对
        while len(self.word_pairs_queue) < batch_size:
            for _ in range(1000):
                # 如果索引等于单词索引列表的长度，即读到头了，则对索引重新赋值
                if self.index == len(self.wordId_list):
                    self.index = 0
                # 中心词索引
                wordId_w = self.wordId_list[self.index]
                # 遍历窗口
                for i in range(max(self.index - window_size, 0),
                               min(self.index + window_size + 1, len(self.wordId_list))):
                    # 上下文单词索引
                    wordId_v = self.wordId_list[i]
                    # 上下文=中心词 跳过
                    if self.index == i:
                        continue
                    # 队列中添加中心词索引与上下文词索引
                    self.word_pairs_queue.append((wordId_w, wordId_v))
                # 中心词索引加1
                self.index += 1
        # 返回mini-batch大小的正采样对
        result_pairs = []
        # 遍历bath大小次
        for _ in range(batch_size):
            # 在result_pairs中添加batch次
            result_pairs.append(self.word_pairs_queue.popleft())
        # 返回mini - batch大小的正采样对
        return result_pairs

    # 获取负采样 输入正采样对数组 positive_pairs，以及每个正采样对需要的负采样数 neg_count 从采样表抽取负采样词的id
    # （假设数据够大，不考虑负采样=正采样的小概率情况）
    def get_negative_sampling(self, positive_pairs, neg_count):
        neg_v = np.random.choice(self.sample_table, size=(len(positive_pairs), neg_count)).tolist()
        return neg_v

    # 估计数据中正采样对数，用于设定batch
    def evaluate_pairs_count(self, window_size):
        return self.word_count_sum * (2 * window_size) - self.sentence_count * (
                1 + window_size) * window_size




if __name__ == "__main__":
    # 读取输入数据
    test_data = InputData('../data/text8.txt', 1)
    # 估计数据中正采样对数，用于设定batch
    test_data.evaluate_pairs_count(2)
    # 返回mini - batch大小的正采样对
    pos_pairs = test_data.get_batch_pairs(10, 2)
    print('正采样:')
    print(pos_pairs)
    pos_word_pairs = []
    # 正采样单词对
    for pair in pos_pairs:
        pos_word_pairs.append((test_data.id2word_dict[pair[0]], test_data.id2word_dict[pair[1]]))
    print(pos_word_pairs)
    # 获取负采样
    neg_pair = test_data.get_negative_sampling(pos_pairs, 3)
    print('负采样:')
    print(neg_pair)
    neg_word_pair = []
    # 负采样单词对
    for pair in neg_pair:
        neg_word_pair.append(
            (test_data.id2word_dict[pair[0]], test_data.id2word_dict[pair[1]], test_data.id2word_dict[pair[2]]))
    print(neg_word_pair)
